import os
import sys
import httpx
from tqdm.auto import tqdm

current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.join(current_dir, ""))
sys.path.append(os.path.join(current_dir, "../"))
sys.path.append(os.path.join(current_dir, "../.."))


def param_type(source_type: str):
    """
    将文档中的数据类型转换成python数据类型
    :param source_type:
    :return:
    """
    if source_type == "str":
        return "str"
    elif source_type == "int":
        return "int"
    elif source_type == "float":
        return "float"
    elif source_type == "datetime":
        return "str"
    else:
        return "str"


def indent(source: str, count: int):
    """
    代码缩进
    :param source:
    :param count:
    :return:
    """
    return "\n".join([(" " * count) + str(row) for row in source])


def columns_dict(return_list: list):
    """
    生成列类型词典
    :param return_list:
    :return:
    """
    col_str = "column_type_dict = {" + (
        ", ".join(["'" + p['name'] + "' : '" + param_type(p['data_type']) + "'" for p in return_list if
                   'name' in p and 'data_type' in p])) + "}"
    return col_str


def fix_column_name(col):
    if col is None:
        return ""
    else:
        return col.replace("'", "").replace("%", "百分比").replace("(", "").replace(")", "").replace("（", '').replace(
            "）",
            "").replace(
            " ", "")


def columns_name_dict(return_list: list):
    """
    生成列中英文名称词典
    :param return_list:
    :return:
    """
    pair_list = [
        "'" + str(p['name']) + "': '" + (
            str(fix_column_name(p['desc'])) if str(p['name']) != 'ts_code' else 'ts_code') + "'"
        for
        p in return_list if 'name' in p and 'desc' in p]

    col_str = "column_name_dict = {" + (
        ", ".join(pair_list)) + "}"
    return col_str


def convert_params(params_list: list):
    """
    生成参数
    :param params_list:
    :return:
    """
    params_dict = {param['name']: param for param in params_list}
    return ", ".join([p['name'] + ": " + param_type(p['data_type']) + " = None" for p in params_dict.values() if
                      p['name'] != 'fields']) + ",fields:str='' "


def all_api(api_dict_list):
    """
    生成all_api 方法,方便接口查询
    :param api_dict_list:
    :return:
    """
    return """
def all_api(keyword:str = None):
    all_api_list={api_dict_list}
    df=pd.DataFrame(all_api_list)
    if keyword is None:
        return df
    else:
        return df[(df['method'].str.contains(keyword))|(df['description'].str.contains(keyword))]        
    """.format(**locals())


def template(method: str, description: str, params: str, param_doc: str, return_doc: str, doc_url: str,
             column_dict: str, column_name_dict: str):
    """
    生成对应的接口函数
    :param module:
    :param method:
    :param description:
    :param params:
    :param param_doc:
    :param return_doc:
    :param doc_url:
    :param column_dict:
    :param column_name_dict:
    :return:
    """

    dict_value = "{}"
    return """
    def {method}(self,{params})->Query:
        \"\"\"
        {description}
        {doc_url}
        <参数>
{param_doc}
        <返回结果>
{return_doc}
        \"\"\"
        params = locals().copy()
        real_params = {dict_value}
        {column_dict}
        {column_name_dict}
        for k, v in params.items():
            if v is not None and  k !='self':
                real_params[k] = v
        return Query(method="{method}",fields=fields, params=real_params,return_column_type_dict=column_type_dict,return_column_name_dict=column_name_dict,token=self.token)    
    """.format(**locals())


segment_pro_bar = """
    def pro_bar(self, ts_code='', start_date='', end_date='', freq='D', asset='E',
                exchange='',
                adj=None,
                ma=[],
                factors=None,
                adjfactor=False,
                retry_count=3):
        \"\"\"
        BAR数据
        Parameters:
        ------------
        ts_code:证券代码，支持股票,ETF/LOF,期货/期权,港股,数字货币
        start_date:开始日期  YYYYMMDD
        end_date:结束日期 YYYYMMDD
        freq:支持日/周/月/季/年
        asset:证券类型 E:股票和交易所基金，I:沪深指数,FT:期货 FD:基金/O期权/H港股/CB可转债
        exchange:市场代码，用户数字货币行情
        adj:复权类型,None不复权,qfq:前复权,hfq:后复权
        factors因子数据，目前支持以下两种：
            vr:量比,默认不返回，返回需指定：factor=['vr']
            tor:换手率，默认不返回，返回需指定：factor=['tor']
                        以上两种都需要：factor=['vr', 'tor']
        retry_count:网络重试次数

        Return
        ----------
        DataFrame
        code:代码
        open：开盘close/high/low/vol成交量/amount成交额/maN均价/vr量比/tor换手率

             期货(asset='FT')
        code/open/close/high/low/avg_price：均价  position：持仓量  vol：成交总量
        \"\"\"
        if (ts_code == '' or ts_code is None) and (adj is not None):
            print('提取复权行情必须输入ts_code参数')
            return
        if len(freq.strip()) >= 3:
            freq = freq.strip().lower()
        else:
            freq = freq.strip().upper() if asset != 'C' else freq.strip().lower()

        if 'min' not in freq:
            today = dt.datetime.today().date()
            today = str(today)[0:10]
            start_date = '' if start_date is None else start_date
            end_date = today if end_date == '' or end_date is None else end_date
            start_date = start_date.replace('-', '')
            end_date = end_date.replace('-', '')
        ts_code = ts_code.strip().upper() if asset != 'C' else ts_code.strip().lower()
        asset = asset.strip().upper()
        for _ in range(retry_count):
            try:
                if asset == 'E':
                    if freq == 'D':
                        data = self.daily(ts_code=ts_code, start_date=start_date, end_date=end_date).get()
                        if factors is not None and len(factors) > 0:
                            ds = self.daily_basic(ts_code=ts_code, start_date=start_date, end_date=end_date).get()[
                                    ['trade_date', 'turnover_rate', 'volume_ratio']]
                            ds = ds.set_index('trade_date')
                            data = data.set_index('trade_date')
                            data = data.merge(ds, left_index=True, right_index=True)
                            data = data.reset_index()
                            if ('tor' in factors) and ('vr' not in factors):
                                data = data.drop('volume_ratio', axis=1)
                            if ('vr' in factors) and ('tor' not in factors):
                                data = data.drop('turnover_rate', axis=1)
                    if freq == 'W':
                        data = self.weekly(ts_code=ts_code, start_date=start_date, end_date=end_date).get()
                    if freq == 'M':
                        data = self.monthly(ts_code=ts_code, start_date=start_date, end_date=end_date).get()
                    if adj is not None:
                        fcts = self.adj_factor(ts_code=ts_code, start_date=start_date, end_date=end_date).get()[
                            ['trade_date', 'adj_factor']]
                        if fcts.shape[0] == 0:
                            return None
                        data = data.set_index('trade_date', drop=False).merge(fcts.set_index('trade_date'),
                                                                              left_index=True,
                                                                              right_index=True, how='left')
                        data['adj_factor'] = data['adj_factor'].bfill().fillna(1)
                        for col in PRICE_COLS:
                            if adj == 'hfq':
                                data[col] = data[col] * data['adj_factor']
                            if adj == 'qfq':
                                data[col] = data[col] * data['adj_factor'] / float(fcts['adj_factor'][0])
                            data[col] = data[col].map(FORMAT)
                            data[col] = data[col].astype(float)
                        if adjfactor is False:
                            data = data.drop('adj_factor', axis=1)
                        if 'min' not in freq:
                            data['change'] = data['close'] - data['pre_close']
                            data['pct_chg'] = data['change'] / data['pre_close'] * 100
                            data['pct_chg'] = data['pct_chg'].map(lambda x: FORMAT(x)).astype(float)
                        else:
                            data = data.drop(['trade_date', 'pre_close'], axis=1)
                elif asset == 'I':
                    if freq == 'D':
                        data = self.index_daily(ts_code=ts_code, start_date=start_date, end_date=end_date).get()
                    if freq == 'W':
                        data = self.index_weekly(ts_code=ts_code, start_date=start_date, end_date=end_date).get()
                    if freq == 'M':
                        data = self.index_monthly(ts_code=ts_code, start_date=start_date, end_date=end_date).get()
                elif asset == 'FT':
                    if freq == 'D':
                        data = self.fut_daily(ts_code=ts_code, start_date=start_date, end_date=end_date,
                                              exchange=exchange).get()
                elif asset == 'O':
                    if freq == 'D':
                        data = self.opt_daily(ts_code=ts_code, start_date=start_date, end_date=end_date,
                                              exchange=exchange).get()
                elif asset == 'CB':
                    if freq == 'D':
                        data = self.cb_daily(ts_code=ts_code, start_date=start_date, end_date=end_date).get()
                elif asset == 'FD':
                    if freq == 'D':
                        data = self.fund_daily(ts_code=ts_code, start_date=start_date, end_date=end_date).get()
                if ma is not None and len(ma) > 0:
                    for a in ma:
                        if isinstance(a, int):
                            data['ma%s' % a] = MA(data['close'], a).map(FORMAT).shift(-(a - 1))
                            data['ma%s' % a] = data['ma%s' % a].astype(float)
                            data['ma_v_%s' % a] = MA(data['vol'], a).map(FORMAT).shift(-(a - 1))
                            data['ma_v_%s' % a] = data['ma_v_%s' % a].astype(float)
                data = data.reset_index(drop=True)
            except Exception as e:
                print(e)
            else:
                return data
        raise IOError('ERROR.')

"""


def extract_doc(doc_list):
    ret_list = []
    for doc in doc_list:
        if doc is not None:
            ret_list.append((doc['name'] if "name" in doc else "") + ":" + (
                doc['data_type'] if "data_type" in doc and doc['data_type'] is not None else "") + " " + (
                                doc['desc'] if "desc" in doc and doc['desc'] is not None else "") + " " + (
                                doc['must'] if "must" in doc and doc['must'] is not None else ""))
    return ret_list


def generate_api():
    headers = {
        'Cookie': 'session-id=015bd125-b0c1-4aa1-be7b-e77d728d18e3; uid=2|1:0|10:1758531860|3:uid|8:MjMyNTE2|ef48d464ff08e84f714efc47fa4c5724adf9dd1a97a26c876c21e0a715a80269; username="2|1:0|10:1758531860|8:username|16:Zm9yICoqKiB0ZXI=|6cd152e87769ce1e8600aaa92e2fb1821d4746a25a50859414ebb807407d9879"'}
    resp = httpx.get("https://tushare.pro/wctapi/documents/tree", headers=headers)
    data = resp.json()['data']

    def flat(data, ret):
        for item in data:
            if len(item['children']) == 0:
                ret.append(item)
            else:
                flat(item['children'], ret)

    ret = []
    flat(data, ret)
    api_list = []
    codes = []
    api_dict_list = []
    methods = []
    for api in tqdm(ret):
        res = httpx.get(f"https://tushare.pro/wctapi/documents/{api['id']}/api_info", headers=headers)
        data = res.json()['data']
        if data is not None:
            params_doc = data['inputs']
            return_doc = data['outputs']
            params_list = [ipt for ipt in data['inputs']]
            return_list = [opt for opt in data['outputs']]
            params = convert_params(params_list)
            column_dict = columns_dict(return_list)
            name_dict = columns_name_dict(return_list)
            method = data['name']
            codes.append(
                template(method, data['title'], params, indent(extract_doc(params_doc), 8),
                         indent(extract_doc(return_doc), 8),
                         f"https://tushare.pro/document/2?doc_id={api['id']}",
                         column_dict, name_dict))
            methods.append(method)
            api_dict_list.append({"method": method, "description": "",
                                  'url': f"https://tushare.pro/document/2?doc_id={api['id']}"})
            api_list.append(data)

    with open("api.py", 'w', encoding='utf-8') as f:
        f.write(
            "from .models import Query\nfrom datetime import datetime\nfrom .utils import store_config, load_config\nimport numpy as np\nimport datetime as dt\nfrom .cache import cache\n")
        f.write("""
import pandas as pd
import datetime as dt

PRICE_COLS = ['open', 'close', 'high', 'low', 'pre_close']
FORMAT = lambda x: '%.4f' % x
FREQS = {'D': '1DAY',
         'W': '1WEEK',
         'Y': '1YEAR',
         }
FACT_LIST = {
    'tor': 'turnover_rate',
    'turnover_rate': 'turnover_rate',
    'vr': 'volume_ratio',
    'volume_ratio': 'volume_ratio',
    'pe': 'pe',
    'pe_ttm': 'pe_ttm',
}
def MA(DF, N):
    return pd.Series.rolling(DF, N).mean()
    
class TushareClient(object):
    def __init__(self, token=None):
        self.config=load_config()
        if token is None or len(token) < 1:
            self.token = self.config.get("token")
            if self.token is None:
                print("本地没有找到token缓存,请输入token!")
        else:
            self.token = token
            self.config['token'] = token
            store_config(self.config)

    def trade_days(self, start_date=None, end_date=None, count=None) -> np.array:
        \"\"\"
        获取指定范围交易日
        :param start_date: 开始日期
        :param end_date: 结束日期
        :param count: 数量
        :return:
        \"\"\"
        end_date = dt.datetime.now() if end_date is None else end_date
        key = "交易日_base"
        with cache:
            cached = cache.get(key)
        if cached is None:
            df = self.trade_cal(end_date=(dt.datetime.now() + dt.timedelta(days=365)).strftime("%Y%m%d"),
                                is_open='1').get()
            days = np.array(pd.to_datetime(df['cal_date']).tolist())
            with cache:
                cache.set(key, days, expire=86400, retry=True)
            cached = days
        if start_date is None:
            return cached[cached <= pd.to_datetime(end_date.strftime("%Y-%m-%d"))][-count:]
        else:
            return cached[(cached >= pd.to_datetime(start_date.strftime("%Y-%m-%d"))) & (
                    cached <= pd.to_datetime(end_date.strftime("%Y-%m-%d")))]

    def last_closed_trade_day(self, end_date=None) -> dt.datetime:
        \"\"\"最后一个收盘的交易日\"\"\"
        end_date = dt.datetime.now() if end_date is None else end_date
        trade_days = self.trade_days(end_date=end_date, count=5)
        if (dt.datetime.now() - trade_days[-1]).days == 0 and end_date.hour * 60 + end_date.minute < 901:
            # 今天是交易日，沒有收盤
            return trade_days[-2]
        else:
            return trade_days[-1]

    @staticmethod
    def last_quarters(count=10, end_date=None):
        \"\"\"
        最新报告期
        :param end_date:
        :param count:
        :return:
        \"\"\"
        end_date = dt.datetime.now() if end_date is None else end_date
        return pd.date_range(end=end_date.date(), periods=count, freq="Q")

    @staticmethod
    def clean_cache(tag: str = "api"):
        \"\"\"
        清理指定类型的api缓存
        :param tag:
        :return:
        \"\"\"
        with cache:
            cache.evict(tag=tag)  
        """)
        f.write("\n\n".join(codes))
        f.write("\n" + segment_pro_bar)
        # f.write("\n\n" + all_api(api_dict_list))
        # f.write("\n\n__all__ = ['" + "', '".join(methods) + "']\n")


if __name__ == "__main__":
    generate_api()
